#!/usr/bin/env python3
"""
通用评估模块使用演示

展示如何在不同场景下使用通用的CLIP检索评估函数
"""

import torch
import numpy as np
from eval import (
    compute_clip_retrieval_metrics,
    print_metrics,
    quick_eval,
    compute_similarity_matrix,
    compute_recall_at_k
)

def demo_basic_usage():
    """演示基本用法"""
    print("=== 基本用法演示 ===")
    
    # 模拟从模型提取的特征
    batch_size = 8
    feature_dim = 512
    
    torch.manual_seed(42)  # 确保可重现
    image_features = torch.randn(batch_size, feature_dim)
    text_features = torch.randn(batch_size, feature_dim)
    
    # 计算检索指标
    metrics = compute_clip_retrieval_metrics(image_features, text_features)
    
    # 打印结果
    print_metrics(metrics, title="随机特征评估结果")
    
    return metrics

def demo_perfect_case():
    """演示完美情况"""
    print("\n=== 完美情况演示 ===")
    
    # 创建理想情况：每个图像与对应文本完全匹配
    batch_size = 5
    feature_dim = 256
    
    # 创建相同的特征对
    base_features = torch.randn(batch_size, feature_dim)
    image_features = base_features + 0.01 * torch.randn_like(base_features)  # 添加微小噪声
    text_features = base_features + 0.01 * torch.randn_like(base_features)
    
    metrics = compute_clip_retrieval_metrics(image_features, text_features)
    print_metrics(metrics, title="理想匹配情况")
    
    return metrics

def demo_difficult_case():
    """演示困难情况"""
    print("\n=== 困难情况演示 ===")
    
    # 手动构造一个困难的相似度矩阵
    n = 4
    similarities = torch.tensor([
        [0.3, 0.7, 0.4, 0.2],  # image 0 最相似 text 1 (错误)
        [0.5, 0.8, 0.3, 0.1],  # image 1 最相似 text 1 (正确)
        [0.2, 0.3, 0.9, 0.4],  # image 2 最相似 text 2 (正确)
        [0.4, 0.2, 0.6, 0.1],  # image 3 最相似 text 2 (错误)
    ])
    
    print("相似度矩阵:")
    print(similarities.numpy())
    print(f"对角线值 (正确匹配): {torch.diag(similarities).tolist()}")
    
    # 手动计算各个召回率
    print("\n手动计算结果:")
    i2t_r1 = compute_recall_at_k(similarities, k=1, direction='i2t')
    t2i_r1 = compute_recall_at_k(similarities, k=1, direction='t2i')
    i2t_r2 = compute_recall_at_k(similarities, k=2, direction='i2t')
    t2i_r2 = compute_recall_at_k(similarities, k=2, direction='t2i')
    
    print(f"I2T R@1: {i2t_r1:.3f} (期望: 0.500)")
    print(f"T2I R@1: {t2i_r1:.3f} (期望: 0.500)")
    print(f"I2T R@2: {i2t_r2:.3f}")
    print(f"T2I R@2: {t2i_r2:.3f}")
    print(f"Mean R@1: {(i2t_r1 + t2i_r1) / 2:.3f} (期望: 0.500)")

def demo_temperature_effect():
    """演示温度参数效果"""
    print("\n=== 温度参数效果演示 ===")
    
    # 创建测试数据
    torch.manual_seed(123)
    image_features = torch.randn(6, 128)
    text_features = torch.randn(6, 128)
    
    temperatures = [0.5, 1.0, 2.0, 4.0]
    
    for temp in temperatures:
        metrics = compute_clip_retrieval_metrics(
            image_features, text_features, 
            temperature=temp, k_values=[1, 5]
        )
        print(f"Temperature {temp}: Mean R@1 = {metrics.mean_r1:.3f}, Mean R@5 = {metrics.mean_r5:.3f}")

def demo_quick_functions():
    """演示快捷函数"""
    print("\n=== 快捷函数演示 ===")
    
    # 使用快速评估函数
    image_features = torch.randn(10, 256)
    text_features = torch.randn(10, 256)
    
    r1, r5, r10 = quick_eval(image_features, text_features)
    print(f"快速评估结果: R@1={r1:.3f}, R@5={r5:.3f}, R@10={r10:.3f}")
    
    # 使用相似度矩阵计算函数
    similarities = compute_similarity_matrix(image_features, text_features)
    print(f"相似度矩阵形状: {similarities.shape}")
    print(f"相似度范围: [{similarities.min():.3f}, {similarities.max():.3f}]")

def demo_integration_example():
    """演示与现有代码集成"""
    print("\n=== 集成示例 ===")
    
    # 模拟一个训练循环中的验证步骤
    class MockModel:
        def validate_step(self, batch):
            # 模拟模型前向传播
            batch_size = 16
            
            # 模拟提取特征
            image_features = torch.randn(batch_size, 512)
            text_features = torch.randn(batch_size, 512) 
            
            # 使用通用评估函数
            metrics = compute_clip_retrieval_metrics(
                image_features, text_features,
                k_values=[1, 5, 10],
                temperature=1.0,
                normalize=True
            )
            
            return metrics
    
    # 模拟使用
    model = MockModel()
    metrics = model.validate_step(None)
    
    print("模拟验证结果:")
    print(f"  I2T R@1: {metrics.i2t_r1:.4f}")
    print(f"  T2I R@1: {metrics.t2i_r1:.4f}")
    print(f"  Mean R@1: {metrics.mean_r1:.4f}")
    
    # 转换为字典用于日志记录
    metrics_dict = metrics.to_dict()
    print(f"日志格式: {list(metrics_dict.keys())[:3]}...")

def main():
    """主演示函数"""
    print("🚀 WenwuClip 通用评估模块演示")
    print("=" * 50)
    
    # 运行各种演示
    demo_basic_usage()
    demo_perfect_case()
    demo_difficult_case()
    demo_temperature_effect()
    demo_quick_functions()
    demo_integration_example()
    
    print("\n" + "=" * 50)
    print("✅ 演示完成！")
    print("\n📋 使用指南:")
    print("1. 导入: from common.eval import compute_clip_retrieval_metrics")
    print("2. 计算: metrics = compute_clip_retrieval_metrics(img_feat, txt_feat)")
    print("3. 使用: metrics.mean_r1, metrics.to_dict()")
    print("4. 打印: print_metrics(metrics)")

if __name__ == "__main__":
    main()